import dns.message
import dns.rrset
import dns.rdata
import socketserver
from base64 import b64encode
import logging

logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')

data_store = {
    "weather": "Sunny, 25°C",
    "news": "Local team wins championship",
    "quote": "The only way to do great work is to love what you do."
}

class DNSHandler(socketserver.BaseRequestHandler):
    def handle(self):
        socket = self.request[1]
        data = self.request[0]

        try:
            msg = dns.message.from_wire(data)
            query = str(msg.question[0].name)
            logging.info(f"Received query: {query}")

            parts = query.split('.')
            key = parts[0]
            response_data = data_store.get(key, "Data not found")
            encoded_data = b64encode(response_data.encode()).decode()

            response = dns.message.make_response(msg)
            rdata = dns.rdata.from_text(dns.rdataclass.IN, dns.rdatatype.TXT, f'"{encoded_data}"')
            rrset = dns.rrset.from_rdata(msg.question[0].name, 300, rdata)
            response.answer.append(rrset)

            logging.info(f"Sending response for key '{key}': {response_data}")
            socket.sendto(response.to_wire(), self.client_address)
        except Exception as e:
            logging.error(f"Error processing request: {str(e)}")
            error_response = dns.message.make_response(msg)
            socket.sendto(error_response.to_wire(), self.client_address)

if __name__ == '__main__':
    server = socketserver.UDPServer(('', 53), DNSHandler)
    logging.info("Starting DNS server...")
    server.serve_forever()
    